import numpy
from scipy.io import wavfile
import argparse
import math
import os
import sys

sys.path.append('audio_feature')
import extract_feature
from  spectrogram  import (
fs,FFTLength,frameSamples,overlapSamples,numBands
)
import spectrogram as spec
import filter_bank
import window


#https://haythamfayek.com/2016/04/21/speech-processing-for-machine-learning.html
#Framing
#ypical frame sizes in speech processing range from 20 ms to 40 ms with 50% (+/-10%)
# overlap between consecutive frames. Popular settings are 25 ms for the frame size,
# frame_size = 0.025 and a 10 ms stride (15 ms overlap), frame_stride = 0.01.
#frame_size = 0.025
#frame_stride = 0.01

#frame_length, frame_step = frame_size * fs, frame_stride * fs  # Convert from seconds to samples

#frame_length = int(round(frame_length))
#frame_step = int(round(frame_step))




#https://haythamfayek.com/2016/04/21/speech-processing-for-machine-learning.html
#Pre-Emphasis
def pre_emphasis(signal):
    pre_emphasis = 0.97
    emphasized_signal = numpy.append(signal[0], signal[1:] - pre_emphasis * signal[:-1])
    return emphasized_signal

def process_file(wav_file_fullpath, out_dir, base_file_name):
    spec_file_name = os.path.join(out_dir, base_file_name.replace(".wav", ".npy", 1))

    sample_rate, samples = wavfile.read(wav_file_fullpath)

    data = samples / 32768
    # pre-emphasis
    #data = pre_emphasis(data)
    spectrogram = spec.gen_spectrogram(data)
    numpy.save(spec_file_name,spectrogram)


parser = argparse.ArgumentParser(description='scan .wav files and gen spectrogram files.')
parser.add_argument("--src", required=True)
parser.add_argument("--dst", default="spec")
parser.add_argument("--exclude", nargs='*', default=[])
args = parser.parse_args()

in_path = args.src
out_path = args.dst

bark_fbank, _, _ = filter_bank.gen_bark_filter_bank(FFTLength, fs, numBands)
spec.afe = extract_feature.audioFeatureExtractor(
    SampleRate=fs,
    FFTLength=FFTLength,
    Window=window.hann(frameSamples, "periodic"),
    OverlapLength=overlapSamples,
    FilterBank=bark_fbank)

if not os.path.exists(out_path):
    os.mkdir(out_path)

if os.path.isfile(in_path):
    file = os.path.basename(in_path)
    process_file(in_path, out_path, file)

    sys.exit(0)
# first level subdir is datastore: test/validation/train
datastore_list=[]
for root, dirs, _ in os.walk(in_path):
    for e in dirs:
        if e not in args.exclude:
            # print(e)
            # print(in_dir)
            datastore_list.append(e)

# walk through each datastore dir
for ds in datastore_list:
    ds_in_path = os.path.join(in_path, ds)
    ds_out_path = os.path.join(out_path, ds)
    if not os.path.exists(ds_out_path):
        os.mkdir(ds_out_path)
    print("%s->%s"%(ds_in_path,ds_out_path))
    for r, dirs, files in os.walk(ds_in_path):
        for e in dirs:
            in_dir = os.path.join(ds_in_path, e)
            out_dir = os.path.join(ds_out_path, e)
            if not os.path.exists(out_dir):
                os.mkdir(out_dir)
            # print(out_dir)
            for r, _, files in os.walk(in_dir):
                for file in files:
                    # print(file)
                    if '.wav' in file:
                        in_fullpath = os.path.join(r, file)
                        process_file(in_fullpath, out_dir, file)
